"""
ConnLLM提供商注册
用于动态注册和发现LLM提供商实现
"""
import os
import sys
import logging
import importlib
import importlib.util
from pathlib import Path
from typing import Dict, Type, Callable, Optional, Set, List, Any

# 设置日志
logger = logging.getLogger("connllm.registry")

# 类型提示
ProviderClass = Type[Any]  # 动态类型，避免循环导入
ProviderDecorator = Callable[[ProviderClass], ProviderClass]


class ProviderRegistry:
	"""提供商注册表，管理可用的LLM提供商实现"""
	
	def __init__(self):
		"""初始化注册表"""
		self._providers: Dict[str, ProviderClass] = {}
		self._providers_paths: List[str] = []
	
	def register(self, name: str, provider_class: ProviderClass) -> None:
		"""
		注册提供商类
		
		Args:
			name: 提供商名称
			provider_class: 提供商类
		"""
		self._providers[name] = provider_class
		logger.debug(f"注册提供商: {name}")
	
	def get(self, name: str) -> Optional[ProviderClass]:
		"""
		获取提供商类
		
		Args:
			name: 提供商名称
			
		Returns:
			提供商类，如果不存在则返回None
		"""
		return self._providers.get(name)
	
	def list_providers(self) -> List[str]:
		"""
		获取所有已注册的提供商名称
		
		Returns:
			提供商名称列表
		"""
		return list(self._providers.keys())
	
	def add_providers_path(self, path: str) -> None:
		"""
		添加提供商模块路径
		
		Args:
			path: 提供商模块路径
		"""
		if path not in self._providers_paths:
			self._providers_paths.append(path)
	
	def get_providers_paths(self) -> List[str]:
		"""
		获取所有提供商模块路径
		
		Returns:
			提供商模块路径列表
		"""
		paths = list(self._providers_paths)
		
		# 如果没有设置路径，添加默认路径
		if not paths:
			# 尝试获取包的提供商路径
			try:
				import connllm
				package_dir = Path(connllm.__file__).parent
				paths.append(str(package_dir / "providers"))
			except ImportError:
				# 如果包未安装，使用相对路径
				base_dir = Path(__file__).parent.parent
				paths.append(str(base_dir / "providers"))
		
		return paths
	
	def discover_providers(self) -> Set[str]:
		"""
		发现并加载所有提供商模块
		
		Returns:
			发现的提供商集合
		"""
		discovered = set()
		
		# 获取所有提供商路径
		paths = self.get_providers_paths()
		
		for path in paths:
			discovered.update(self._discover_from_path(path))
		
		return discovered
	
	def _discover_from_path(self, path: str) -> Set[str]:
		"""
		从指定路径发现并加载提供商模块
		
		Args:
			path: 提供商模块路径
			
		Returns:
			发现的提供商集合
		"""
		discovered = set()
		
		# 确保路径存在
		if not os.path.exists(path) or not os.path.isdir(path):
			logger.debug(f"提供商路径不存在: {path}")
			return discovered
		
		# 将路径添加到sys.path，以便能够导入模块
		if path not in sys.path:
			sys.path.insert(0, os.path.dirname(path))
		
		# 获取目录名称
		dir_name = os.path.basename(path)
		
		# 遍历目录中的Python文件
		for filename in os.listdir(path):
			if filename.startswith("_") or not filename.endswith(".py"):
				continue
			
			module_name = filename[:-3]  # 去掉.py后缀
			full_module_name = f"connllm.{dir_name}.{module_name}"
			
			try:
				# 尝试多种导入方式
				try:
					# 直接导入
					module = importlib.import_module(full_module_name)
				except ImportError:
					# 尝试使用相对导入
					rel_module_name = f".{dir_name}.{module_name}"
					try:
						module = importlib.import_module(rel_module_name, package="connllm")
					except ImportError:
						# 尝试使用spec导入
						spec = importlib.util.spec_from_file_location(
							module_name, os.path.join(path, filename)
						)
						if spec is None or spec.loader is None:
							raise ImportError(f"无法从文件加载模块: {filename}")
							
						module = importlib.util.module_from_spec(spec)
						spec.loader.exec_module(module)
				
				# 检查模块是否定义了提供商类
				if module_name in self._providers:
					discovered.add(module_name)
					logger.debug(f"已发现提供商模块: {module_name}")
			except (ImportError, AttributeError) as e:
				logger.warning(f"导入提供商模块 {module_name} 失败: {str(e)}")
			except Exception as e:
				logger.error(f"导入提供商模块 {module_name} 时发生未知错误: {str(e)}")
		
		return discovered
	
	def import_provider(self, provider_type: str) -> bool:
		"""
		导入特定提供商模块
		
		Args:
			provider_type: 提供商类型
			
		Returns:
			是否成功导入
		"""
		# 尝试所有提供商路径
		for path in self.get_providers_paths():
			provider_file = os.path.join(path, f"{provider_type}.py")
			
			if os.path.exists(provider_file):
				try:
					dir_name = os.path.basename(path)
					# 尝试多种导入方式
					try:
						# 直接导入
						importlib.import_module(f"connllm.{dir_name}.{provider_type}")
					except ImportError:
						# 尝试使用相对导入
						try:
							importlib.import_module(f".{dir_name}.{provider_type}", package="connllm")
						except ImportError:
							# 尝试使用spec导入
							spec = importlib.util.spec_from_file_location(
								provider_type, provider_file
							)
							if spec is None or spec.loader is None:
								raise ImportError(f"无法从文件加载模块: {provider_file}")
								
							module = importlib.util.module_from_spec(spec)
							spec.loader.exec_module(module)
					
					return provider_type in self._providers
				except ImportError as e:
					logger.warning(f"导入提供商 {provider_type} 失败: {str(e)}")
				except Exception as e:
					logger.error(f"导入提供商 {provider_type} 时发生未知错误: {str(e)}")
		
		# 尝试使用默认导入路径
		try:
			importlib.import_module(f"connllm.providers.{provider_type}")
			return provider_type in self._providers
		except ImportError:
			pass
		
		logger.warning(f"在所有路径中均未找到提供商 {provider_type}")
		return False
	
	def register_builtin_providers(self) -> None:
		"""注册内置提供商"""
		# 内置提供商模块映射
		BUILTIN_PROVIDERS = {
			"anthropic": "connllm.providers.anthropic",
			"openai": "connllm.providers.openai",
			"openrouter": "connllm.providers.openrouter",
			"moonshot": "connllm.providers.moonshot",
			"ollama": "connllm.providers.ollama",
			"vllm": "connllm.providers.vllm",
		}
		
		for provider, module_path in BUILTIN_PROVIDERS.items():
			try:
				importlib.import_module(module_path)
				logger.debug(f"已注册内置提供商: {provider}")
			except ImportError:
				logger.debug(f"未能加载内置提供商 {provider}: 模块不可用")


# 全局注册表实例
_registry = ProviderRegistry()


def get_registry() -> ProviderRegistry:
	"""
	获取全局注册表实例
	
	Returns:
		注册表实例
	"""
	return _registry


def register_provider(name: str) -> ProviderDecorator:
	"""
	提供商注册装饰器
	
	Args:
		name: 提供商名称
		
	Returns:
		装饰器函数
	"""
	def decorator(cls: ProviderClass) -> ProviderClass:
		_registry.register(name, cls)
		return cls
	return decorator


# 初始化：注册内置提供商
_registry.register_builtin_providers()
